#!/bin/bash

# 检查是否安装了 kubectl
if ! command -v kubectl &> /dev/null; then
    echo "Error: kubectl is not installed. Please install it first."
    exit 1
fi

# 定义文件路径
NAMESPACE_FILE="name"
OUTPUT_FILE="image_list.txt"

# 函数：检查 namespace 文件是否存在
check_namespace_file() {
    if [ ! -f "$NAMESPACE_FILE" ]; then
        echo "Error: $NAMESPACE_FILE file not found in current directory."
        exit 1
    fi
}

# 函数：列出并保存每个 namespace 的镜像
list_images() {
    check_namespace_file
    echo "Saving namespace and image information to $OUTPUT_FILE..."
    echo "NAMESPACE,DEPLOYMENT,IMAGE" > "$OUTPUT_FILE"
    
    while read -r ns; do
        # 跳过空行
        if [ -z "$ns" ]; then
            continue
        fi
        # 获取该 namespace 下所有 deployment 的镜像
        kubectl get deployments -n "$ns" -o custom-columns=":metadata.name,:spec.template.spec.containers[*].image" --no-headers 2>/dev/null | while read -r deployment images; do
            # 检查是否存在多个镜像
            image_count=$(echo "$images" | tr ',' '\n' | wc -l)
            if [ "$image_count" -gt 1 ]; then
                echo "Warning: Namespace $ns, Deployment $deployment has multiple images ($images). Skipping in list."
                continue
            fi
            # 保存 namespace, deployment 和镜像
            echo "$ns,$deployment,$images" >> "$OUTPUT_FILE"
        done
    done < "$NAMESPACE_FILE"
    echo "Image list saved to $OUTPUT_FILE"
}

# 函数：还原镜像
restore_images() {
    check_namespace_file
    if [ ! -f "$OUTPUT_FILE" ]; then
        echo "Error: $OUTPUT_FILE file not found. Run 'image_list -l' first to generate it."
        exit 1
    fi

    echo "Restoring images from $OUTPUT_FILE..."
    # 跳过标题行
    tail -n +2 "$OUTPUT_FILE" | while IFS=',' read -r ns deployment image; do
        # 跳过空行
        if [ -z "$ns" ] || [ -z "$deployment" ] || [ -z "$image" ]; then
            continue
        fi
        # 检查当前 deployment 的镜像数量
        current_images=$(kubectl get deployment "$deployment" -n "$ns" -o custom-columns=":spec.template.spec.containers[*].image" --no-headers 2>/dev/null)
        image_count=$(echo "$current_images" | tr ',' '\n' | wc -l)
        if [ "$image_count" -gt 1 ]; then
            echo "Error: Namespace $ns, Deployment $deployment has multiple images ($current_images). Skipping update."
            exit 1
        fi
        # 还原镜像
        echo "Restoring image for Namespace: $ns, Deployment: $deployment to $image"
        kubectl set image deployment/"$deployment" "*=$image" -n "$ns" 2>/dev/null
        if [ $? -eq 0 ]; then
            echo "Successfully restored $deployment in $ns to $image"
        else
            echo "Error: Failed to restore $deployment in $ns to $image"
            exit 1
        fi
    done
}

# 主逻辑：解析命令行参数
case "$1" in
    -l)
        list_images
        ;;
    -r)
        restore_images
        ;;
    *)
        echo "Usage: $0 -l | -r"
        echo "  -l : List and save namespaces and their corresponding images to $OUTPUT_FILE"
        echo "  -r : Restore images for each namespace from $OUTPUT_FILE"
        exit 1
        ;;
esac